#include "primer/trie.h"
#include <cstddef>
#include <functional>
#include <memory>
#include <string_view>
#include <type_traits>
#include "common/exception.h"

namespace bustub {

template <class T>
auto Trie::Get(std::string_view key) const -> const T * {
  // You should walk through the trie to find the node corresponding to the key. If the node doesn't exist, return
  // nullptr. After you find the node, you should use `dynamic_cast` to cast it to `const TrieNodeWithValue<T> *`. If
  // dynamic_cast returns `nullptr`, it means the type of the value is mismatched, and you should return nullptr.
  // Otherwise, return the value.

  auto node_ptr = root_.get();
  for (const char &c : key) {
    if (node_ptr && node_ptr->children_.count(c)) {
      node_ptr = node_ptr->children_.at(c).get();
    } else {
      return nullptr;
    }
  }
  auto node = dynamic_cast<const TrieNodeWithValue<T> *>(node_ptr);
  if (!node) {
    return nullptr;
  }
  return node->value_.get();
}

template <class T>
auto Trie::Put(std::string_view key, T value) const -> Trie {
  // Note that `T` might be a non-copyable type. Always use `std::move` when creating `shared_ptr` on that value.

  // You should walk through the trie and create new nodes if necessary. If the node corresponding to the key already
  // exists, you should create a new `TrieNodeWithValue`.

  auto node_ptr = root_;
  std::vector<std::pair<std::shared_ptr<const TrieNode>, char>> path;
  int n = key.size();
  for (int i = 0; i < n; ++i) {
    const char &c = key[i];
    path.emplace_back(node_ptr, c);
    if (node_ptr && node_ptr->children_.count(c)) {
      node_ptr = node_ptr->children_.at(c);
    } else {
      node_ptr = nullptr;
    }
  }
  std::shared_ptr<TrieNode> child_ptr = nullptr;
  auto shared_value = std::make_shared<T>(std::move(value));
  // value
  if (node_ptr == nullptr) {
    child_ptr = std::make_shared<TrieNodeWithValue<T>>(shared_value);
  } else {
    child_ptr = std::make_shared<TrieNodeWithValue<T>>(node_ptr->children_, shared_value);
  }
  int pn = path.size();
  for (int j = pn - 1; j >= 0; --j) {
    const auto &[fa, c] = path[j];
    if (fa) {  // clone node
      auto fa_new = fa->Clone();
      fa_new->children_[c] = std::move(child_ptr);
      child_ptr = std::shared_ptr<TrieNode>(std::move(fa_new));
    } else {  // create new node
      auto node =
          std::make_shared<TrieNode>(std::map<char, std::shared_ptr<const TrieNode>>{{c, std::move(child_ptr)}});
      child_ptr = node;
    }
  }
  return Trie(child_ptr);
}

auto Trie::Remove(std::string_view key) const -> Trie {
  // throw NotImplementedException("Trie::Remove is not implemented.");

  // You should walk through the trie and remove nodes if necessary. If the node doesn't contain a value any more,
  // you should convert it to `TrieNode`. If a node doesn't have children any more, you should remove it.

  auto node_ptr = root_;
  std::vector<std::pair<std::shared_ptr<const TrieNode>, char>> path;
  int n = key.size();
  for (int i = 0; i < n; ++i) {
    const char &c = key[i];
    path.emplace_back(node_ptr, c);
    if (node_ptr && (node_ptr->children_.count(c) != 0U)) {
      node_ptr = node_ptr->children_.at(c);
    } else {
      return *this;
    }
  }
  if (!node_ptr->is_value_node_) {
    return *this;
  }
  std::shared_ptr<TrieNode> child_ptr = nullptr;
  if (!node_ptr->children_.empty()) {
    child_ptr = std::make_shared<TrieNode>(node_ptr->children_);
  }
  int pn = path.size();
  for (int j = pn - 1; j >= 0; --j) {
    const auto &[fa, c] = path[j];
    auto fa_new = fa->Clone();
    if (child_ptr == nullptr) {
      fa_new->children_.erase(c);
      if (!fa_new->is_value_node_ && fa_new->children_.empty()) {
        fa_new = nullptr;
      }
    } else {
      fa_new->children_[c] = std::move(child_ptr);
    }
    child_ptr = std::shared_ptr<TrieNode>(std::move(fa_new));
  }
  return Trie(child_ptr);
}

// Below are explicit instantiation of template functions.
//
// Generally people would write the implementation of template classes and functions in the header file. However, we
// separate the implementation into a .cpp file to make things clearer. In order to make the compiler know the
// implementation of the template functions, we need to explicitly instantiate them here, so that they can be picked
// up by the linker.

template auto Trie::Put(std::string_view key, uint32_t value) const -> Trie;
template auto Trie::Get(std::string_view key) const -> const uint32_t *;

template auto Trie::Put(std::string_view key, uint64_t value) const -> Trie;
template auto Trie::Get(std::string_view key) const -> const uint64_t *;

template auto Trie::Put(std::string_view key, std::string value) const -> Trie;
template auto Trie::Get(std::string_view key) const -> const std::string *;

// If your solution cannot compile for non-copy tests, you can remove the below lines to get partial score.

using Integer = std::unique_ptr<uint32_t>;

template auto Trie::Put(std::string_view key, Integer value) const -> Trie;
template auto Trie::Get(std::string_view key) const -> const Integer *;

template auto Trie::Put(std::string_view key, MoveBlocked value) const -> Trie;
template auto Trie::Get(std::string_view key) const -> const MoveBlocked *;

}  // namespace bustub
